#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from importlib import reload
import matplotlib.pyplot as plt
import numpy as np

from kNNAlgo import kNN

# 数据准备
datingDataMat,datingLabels = kNN.file2matrix('datingTestSet.txt')
# 数据输出测试
print(datingDataMat)
print(datingLabels[0:20])
print()

# 数据颜色分类绘制数据
datingDataMatProp1 = np.array([0., 0., 0.])  # [0., 0., 0.]效果相同，但是初始化的时候数据类型不同，np.array([])为{ndarray}，[]为{list}
datingDataMatProp2 = np.array([0., 0., 0.])
datingDataMatProp3 = np.array([0., 0., 0.])

index = 0

# matplotlib 图形化
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1) #441-第一个数为行，第二个为列，第三个为次序


#分颜色绘制数据分类
for colorLabel in datingLabels:
    if colorLabel == 'largeDoses': # 非常喜欢
        datingDataMatProp1 = np.row_stack((datingDataMatProp1, datingDataMat[index])) #矩阵中添加一行
    elif colorLabel == 'smallDoses': # 一般喜欢
        datingDataMatProp2 = np.row_stack((datingDataMatProp2, datingDataMat[index]))
    else:
        if colorLabel == 'didntLike': # 不喜欢
            datingDataMatProp3 = np.row_stack((datingDataMatProp3, datingDataMat[index]))
        else:
            print('datingLabels file error: there is wrong data')
    index += 1

# 数据处理，去除第一行额外的数据
datingDataMatProp1 = np.delete(datingDataMatProp1, 0, 0)
datingDataMatProp2 = np.delete(datingDataMatProp2, 0, 0)
datingDataMatProp3 = np.delete(datingDataMatProp3, 0, 0)

# 绘图 [:, 0]:飞行时长，[:, 1]:玩游戏时间百分比，[:, 2]:吃冰淇淋公升数
type1 = ax.scatter(datingDataMatProp1[:, 0], datingDataMatProp1[:, 1], c='red', s=30)
type2 = ax.scatter(datingDataMatProp2[:, 0], datingDataMatProp2[:, 1], c='blue', s=20)
type3 = ax.scatter(datingDataMatProp3[:, 0], datingDataMatProp3[:, 1], c='black', s=15)

# 加标签
ax.legend((type1, type2, type3), ('极具魅力', '魅力一般', '不喜欢'), loc=0)
plt.xlabel('获取的飞行里程数')
plt.ylabel('玩视频游戏所耗时间百分比')
plt.title('约会配对魅力指数图')
plt.legend()

# 显示图形
plt.show()
